package com.jd.genie.controller;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.TypeReference;
import com.jd.genie.agent.agent.AgentContext;
import com.jd.genie.agent.printer.Printer;
import com.jd.genie.agent.tool.ToolCollection;
import com.jd.genie.agent.tool.common.CodeInterpreterTool;
import com.jd.genie.agent.tool.common.DeepSearchTool;
import com.jd.genie.agent.tool.common.FileTool;
import com.jd.genie.agent.tool.common.ReportTool;
import com.jd.genie.agent.util.DateUtil;
import com.jd.genie.config.GenieConfig;
import com.jd.genie.handler.AgentHandlerFactory;
import com.jd.genie.handler.AgentHandlerService;
import com.jd.genie.model.req.AgentRequest;
import com.jd.genie.model.req.GptQueryReq;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

@Slf4j
@RestController
@RequestMapping("/")
public class GenieController {

    private final ScheduledExecutorService executor = Executors.newScheduledThreadPool(5);

    private static final long HEARTBEAT_INTERVAL = 10_000L; // 10秒心跳间隔

    @Autowired
    protected GenieConfig genieConfig;

    @Autowired
    private AgentHandlerFactory agentHandlerFactory;

    @Autowired
    private IGptProcessService gptProcessService;

    /**
     * 开启SSE心跳
     */
    private ScheduledFuture<?> startHeartbeat(SseEmitter emitter, String requestId)
    {
        return executor.scheduleAtFixedRate(() -> {
            try {
                // 发送心跳消息
                log.info("{} send heartbeat", requestId);
                emitter.send("heartbeat");
            } catch (Exception e) {
                // 发送心跳失败，关闭连接
                log.error("{} heartbeat failed, closing connection", requestId, e);
                emitter.completeWithError(e);
            }
        }, HEARTBEAT_INTERVAL, HEARTBEAT_INTERVAL, TimeUnit.MILLISECONDS);
    }

    /**
     * 注册SSE事件
     */
    private void registerSSEMonitor(SseEmitter emitter, String requestId, ScheduledFuture<?> heartbeatFuture)
    {
        // 监听SSE异常事件
        emitter.onCompletion(() -> {
            log.info("{} SSE connection completed normally", requestId);
            heartbeatFuture.cancel(true);
        });

        // 监听连接超时事件
        emitter.onTimeout(() -> {
            log.info("{} SSE connection timed out", requestId);
            heartbeatFuture.cancel(true);
            emitter.complete();
        });

        // 监听连接错误事件
        emitter.onError((ex) -> {
            log.info("{} SSE connection error: ", requestId, ex);
            heartbeatFuture.cancel(true);
            emitter.completeWithError(ex);
        });
    }

    private static final Map<String, String> outputStylePrompts;

    static {
        String jsonStr = """
                {
                    "html": "",
                    "docs": "，最后以 markdown 展示最终结果",
                    "table": "，最后以excel 展示最终结果",
                    "ppt": "，最后以 ppt 展示最终结果"
                }
                """;
        outputStylePrompts = JSON.parseObject(jsonStr, new TypeReference<>() {
        });
    }

    /**
     * html模式： query+以 html展示
     * docs模式：query+以 markdown展示
     * table 模式: query+以 excel 展示
     */
    public static String handleOutputStyle(AgentRequest request)
    {
        String query = request.getQuery();
        Map<String, String> outputStyleMap = outputStylePrompts;
        if (!StringUtils.isEmpty(request.getOutputStyle())) {
            query += outputStyleMap.computeIfAbsent(request.getOutputStyle(), k -> "");
        }
        return query;
    }


    /**
     * 构建工具列表
     */
    public static ToolCollection buildToolCollection(AgentContext agentContext, Printer printer)
    {
        ToolCollection toolCollection = new ToolCollection(agentContext.getRequestId());
        toolCollection.addTool(new FileTool(agentContext, printer));
        toolCollection.addTool(new CodeInterpreterTool(agentContext, printer));
        toolCollection.addTool(new ReportTool(agentContext, printer));
        toolCollection.addTool(new DeepSearchTool(agentContext, printer));
        return toolCollection;
    }

    /**
     * 执行智能体调度
     */
    @PostMapping("/AutoAgent")
    public SseEmitter AutoAgent(@RequestBody AgentRequest request)
    {
        log.info("{} auto agent request: {}", request.getRequestId(), JSON.toJSONString(request));

        SseEmitter emitter = new SseEmitter(60 * 60 * 1000L);
        // SSE心跳
        ScheduledFuture<?> heartbeatFuture = startHeartbeat(emitter, request.getRequestId());
        // 监听SSE事件
        registerSSEMonitor(emitter, request.getRequestId(), heartbeatFuture);
        // 拼接输出类型
        request.setQuery(handleOutputStyle(request));
        // 执行调度引擎
        new Thread(() -> {
            try {
                Printer printer = new Printer(emitter, request, request.getAgentType());
                AgentContext agentContext = AgentContext.builder()
                        .requestId(request.getRequestId())
                        .sessionId(request.getRequestId())
                        .query(request.getQuery())
                        .task("")
                        .dateInfo(DateUtil.CurrentDateInfo())
                        .productFiles(new ArrayList<>())
                        .taskProductFiles(new ArrayList<>())
                        .agentType(request.getAgentType())
                        .isStream(Objects.nonNull(request.getIsStream()) ? request.getIsStream() : false)
                        .build();

                agentContext.setToolCollection(buildToolCollection(agentContext, printer));
                AgentHandlerService handler = agentHandlerFactory.getHandler(request.getAgentType());
                handler.handle(agentContext, request, printer);
                emitter.complete();
            } catch (Exception e) {
                log.error("{} auto agent error", request.getRequestId(), e);
            }
        }).start();
        return emitter;
    }


    /**
     * 探活接口
     */
    @RequestMapping(value = "/web/health", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public ResponseEntity<String> health()
    {
        return ResponseEntity.ok("ok");
    }


    /**
     * 处理Agent流式增量查询请求，返回SSE事件流
     *
     * @param params 查询请求参数对象，包含GPT查询所需信息
     * @return 返回SSE事件发射器，用于流式传输增量响应结果
     */
    @RequestMapping(value = "/web/api/v1/gpt/queryAgentStreamIncr", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter queryAgentStreamIncr(@RequestBody GptQueryReq params)
    {
        return gptProcessService.queryMultiAgentIncrStream(params);
    }

}
    